package com.mfn.sportjava.knowledge.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.google.common.collect.Lists;
import com.mfn.sportjava.knowledge.model.KnowledgeEntry;
import com.mfn.sportjava.knowledge.mapper.KnowledgeEntryMapper;
import com.mfn.sportjava.knowledge.service.KnowledgePersistenceService;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.MutationResult;
import io.milvus.param.R;
import io.milvus.param.dml.UpsertParam;
import lombok.Getter;
import org.springframework.transaction.annotation.Transactional;
import lombok.RequiredArgsConstructor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 将知识条目持久化到 MySQL和 Milvus。
 *
 * @author Anyang
 * @version 1.0.2
 * @date 2025/05/18
 */
@Service
@RequiredArgsConstructor
public class KnowledgePersistenceServiceImpl implements KnowledgePersistenceService {

    private static final Logger log = LoggerFactory.getLogger(KnowledgePersistenceServiceImpl.class);

    private final KnowledgeEntryMapper knowledgeEntryMapper;
    private final Optional<MilvusServiceClient> milvusServiceClientOpt;
    @NonNull
    private final EmbeddingClient embeddingClient;

    @Value("${spring.milvus.collection-name:match_knowledge}")
    private String collectionName;

    private static final int EMBEDDING_BATCH_SIZE = 10; // Embedding 批处理大小
    private static final String MILVUS_PK_FIELD_NAME = "mysql_id"; // Milvus 中对应 MySQL 主键的字段名
    private static final String VECTOR_FIELD_NAME = "embedding"; // Milvus 中向量字段名
    private static final String TEXT_CONTENT_FIELD_NAME = "text_content"; // Milvus 中存储原文的字段名


    /**
     * 将知识条目保存到 MySQL 和 Milvus。
     *
     * @param entry 要保存的条目
     * @return 保存后的条目
     */
    @Override
    @Transactional(rollbackFor = Exception.class)
    public KnowledgeEntry saveOrUpdateKnowledge(KnowledgeEntry entry) {
        log.info("持久化服务 (单个): 准备保存/更新知识条目，Key: {}", entry.getMatchKey());
        validateKnowledgeEntry(entry, -1); // -1 表示单个条目的上下文

        KnowledgeEntry persistedEntry = persistToMySql(entry);

        if (persistedEntry.getId() == null) {
            log.error("持久化服务 (单个): MySQL保存/更新后条目 ID 仍然为 null, MatchKey: {}. 操作失败。", persistedEntry.getMatchKey());
            throw new IllegalStateException("保存/更新知识条目后未能获取ID: " + persistedEntry.getMatchKey());
        }

        log.info("持久化服务 (单个): MySQL 操作完成, ID: {}", persistedEntry.getId());
        processMilvusPersistence(Collections.singletonList(persistedEntry));
        return persistedEntry;
    }

    /**
     * 将知识条目保存到 MySQL。
     *
     * @param entries 要保存的条目列表
     * @return 保存后的条目列表
     */
    @Override
    @Transactional(rollbackFor = Exception.class)
    public List<KnowledgeEntry> saveOrUpdateKnowledgeBatch(List<KnowledgeEntry> entries) {
        if (CollectionUtils.isEmpty(entries)) {
            log.warn("持久化服务 (批量): 输入的条目列表为空。");
            return Collections.emptyList();
        }
        log.info("持久化服务 (批量): 准备保存/更新 {} 个知识条目。", entries.size());
        validateKnowledgeEntries(entries);

        List<KnowledgeEntry> persistedEntries = persistToMySqlBatch(entries);
        for (KnowledgeEntry entry : persistedEntries) {
            if (entry.getId() == null) {
                log.error("持久化服务 (批量): MySQL批量操作后，条目 (MatchKey: {}) 的ID为null。", entry.getMatchKey());
                throw new IllegalStateException("批量保存/更新知识条目后未能获取部分条目的ID。");
            }
        }

        log.info("持久化服务 (批量): MySQL 操作完成, 共处理 {} 个条目。", persistedEntries.size());
        processMilvusPersistence(persistedEntries);
        return persistedEntries;
    }


    /**
     * 将知识条目保存到 MySQL。
     *
     * @param entry 要保存的条目
     * @return 保存后的条目
     */
    private KnowledgeEntry persistToMySql(KnowledgeEntry entry) {
        QueryWrapper<KnowledgeEntry> queryWrapper = new QueryWrapper<KnowledgeEntry>()
                .eq("match_key", entry.getMatchKey());
        KnowledgeEntry existingEntry = knowledgeEntryMapper.selectOne(queryWrapper);

        if (existingEntry != null) {
            log.info("持久化服务 (MySQL 单个): 找到 MatchKey '{}' 的现有条目 (ID: {}), 进行更新。", entry.getMatchKey(), existingEntry.getId());
            updateExistingMySqlEntry(existingEntry, entry);
            return existingEntry;
        } else {
            log.info("持久化服务 (MySQL 单个): 未找到 MatchKey '{}' 的现有条目, 创建新条目。", entry.getMatchKey());
            knowledgeEntryMapper.insert(entry);
            return entry;
        }
    }

    /**
     * 更新现有的 MySQL 条目。
     *
     * @param entries 要更新的条目列表
     */
    private List<KnowledgeEntry> persistToMySqlBatch(List<KnowledgeEntry> entries) {
        List<String> matchKeys = entries.stream().map(KnowledgeEntry::getMatchKey).distinct().collect(Collectors.toList());
        Map<String, KnowledgeEntry> existingEntriesMap = Collections.emptyMap();
        if (!matchKeys.isEmpty()) {
            existingEntriesMap = knowledgeEntryMapper.findByMatchKeyIn(matchKeys)
                    .stream()
                    .collect(Collectors.toMap(KnowledgeEntry::getMatchKey, Function.identity()));
        }
        log.debug("持久化服务 (MySQL 批量): 从数据库中找到 {} 个已存在的 MatchKeys (共查询 {} 个唯一Key)。", existingEntriesMap.size(), matchKeys.size());

        List<KnowledgeEntry> processedEntries = new ArrayList<>();
        for (KnowledgeEntry newEntry : entries) {
            KnowledgeEntry existingEntry = existingEntriesMap.get(newEntry.getMatchKey());
            if (existingEntry != null) {
                log.debug("持久化服务 (MySQL 批量): 准备更新 MatchKey: {}", newEntry.getMatchKey());
                updateExistingMySqlEntry(existingEntry, newEntry);
                processedEntries.add(existingEntry);
            } else {
                log.debug("持久化服务 (MySQL 批量): 准备插入新 MatchKey: {}", newEntry.getMatchKey());
                knowledgeEntryMapper.insert(newEntry);
                processedEntries.add(newEntry); // ID 应该由 MyBatis-Plus 填充
            }
        }
        return processedEntries;
    }

    /**
     * 更新已有的 MySQL 条目
     * @param existingEntry 待更新的已有条目
     * @param newEntryData 新的条目数据
     */
    private void updateExistingMySqlEntry(KnowledgeEntry existingEntry, KnowledgeEntry newEntryData) {
        existingEntry.setLeague(newEntryData.getLeague());
        existingEntry.setHomeTeam(newEntryData.getHomeTeam());
        existingEntry.setAwayTeam(newEntryData.getAwayTeam());
        existingEntry.setContent(newEntryData.getContent());
        knowledgeEntryMapper.updateById(existingEntry);
    }

    /**
     * 处理 Milvus 持久化逻辑
     * @param entriesToProcess 需要处理的条目列表
     */
    private void processMilvusPersistence(List<KnowledgeEntry> entriesToProcess) {
        if (milvusServiceClientOpt.isEmpty()) {
            log.warn("持久化服务 (Milvus): MilvusServiceClient 不可用。跳过 {} 个条目的 Milvus 操作。", entriesToProcess.size());
            return;
        }
        if (CollectionUtils.isEmpty(entriesToProcess)) {
            log.debug("持久化服务 (Milvus): 没有需要处理的条目。");
            return;
        }

        List<KnowledgeEntry> validEntriesForMilvus = filterValidEntriesForMilvus(entriesToProcess);
        if (validEntriesForMilvus.isEmpty()) {
            log.info("持久化服务 (Milvus): 筛选后没有有效的条目需要进行 Milvus 持久化。");
            return;
        }

        List<MilvusDataTuple> milvusDataTuples = generateEmbeddingsInBatches(validEntriesForMilvus);
        if (milvusDataTuples.isEmpty()) {
            log.warn("持久化服务 (Milvus): Embedding 处理后，没有成功的向量可供 Upsert 到 Milvus。");
            return;
        }

        upsertDataToMilvus(milvusDataTuples, milvusServiceClientOpt.get());
    }

    /**
     * 筛选出有效的条目，并生成 MilvusDataTuple 列表，用于进行 Milvus 持久化。
     *
     * @param entries 待处理的条目列表
     * @return MilvusDataTuple 列表，包含每个条目的 ID 和 Embedding
     */
    private List<KnowledgeEntry> filterValidEntriesForMilvus(List<KnowledgeEntry> entries) {
        return entries.stream()
                .filter(entry -> {
                    if (entry.getId() == null) {
                        log.warn("持久化服务 (Milvus 筛选): 条目 (MatchKey: {}) 缺少 ID，无法进行 Milvus 持久化。", entry.getMatchKey());
                        return false;
                    }
                    if (entry.getContent() == null || entry.getContent().isBlank()) {
                        log.warn("持久化服务 (Milvus 筛选): 条目 (ID: {}) 内容为空，无法进行 Milvus 持久化。", entry.getId());
                        return false;
                    }
                    return true;
                })
                .collect(Collectors.toList());
    }

    /**
     * 为 Milvus 持久化生成 Embedding。
     *
     * @param validEntries 待处理的有效条目
     * @return MilvusDataTuple 列表，包含每个条目的 ID 和 Embedding
     */
    private List<MilvusDataTuple> generateEmbeddingsInBatches(List<KnowledgeEntry> validEntries) {
        log.info("持久化服务 (Embedding): 准备为 {} 个有效条目进行 Embedding (分批大小: {})。",
                validEntries.size(), EMBEDDING_BATCH_SIZE);

        List<MilvusDataTuple> allMilvusData = new ArrayList<>();
        List<List<KnowledgeEntry>> entryPartitions = Lists.partition(validEntries, EMBEDDING_BATCH_SIZE);

        for (List<KnowledgeEntry> batchEntries : entryPartitions) {
            List<String> batchContents = batchEntries.stream()
                    .map(KnowledgeEntry::getContent)
                    .collect(Collectors.toList());
            log.debug("持久化服务 (Embedding): 处理 Embedding 批次，大小: {}", batchContents.size());

            try {
                EmbeddingResponse embeddingResponse = embeddingClient.call(new EmbeddingRequest(batchContents, null));
                if (embeddingResponse != null && !CollectionUtils.isEmpty(embeddingResponse.getResults())) {
                    List<Embedding> batchEmbeddings = embeddingResponse.getResults();
                    if (batchEmbeddings.size() == batchEntries.size()) {
                        for (int i = 0; i < batchEntries.size(); i++) {
                            KnowledgeEntry entry = batchEntries.get(i);
                            Embedding embedding = batchEmbeddings.get(i);
                            if (embedding != null && !CollectionUtils.isEmpty(embedding.getOutput())) {
                                List<Float> floatVector = embedding.getOutput().stream()
                                        .map(Double::floatValue)
                                        .collect(Collectors.toList());
                                allMilvusData.add(new MilvusDataTuple(entry.getId(), floatVector, entry.getContent()));
                            } else {
                                log.warn("持久化服务 (Embedding): 条目 ID {} (MatchKey: {}) 的 Embedding 结果为空或无效，跳过。",
                                        entry.getId(), entry.getMatchKey());
                            }
                        }
                    } else {
                        log.error("持久化服务 (Embedding): 批次 Embedding 返回结果数量 ({}) 与请求数量 ({}) 不匹配，跳过此批次。",
                                batchEmbeddings.size(), batchEntries.size());
                    }
                } else {
                    log.error("持久化服务 (Embedding): 批次 Embedding 调用返回空响应或空结果，跳过此批次。请求内容数量: {}", batchContents.size());
                }
            } catch (Exception e) {
                log.error("持久化服务 (Embedding): 处理 Embedding 批次时发生错误: {}", e.getMessage(), e);
                // 根据需求，可能需要重新抛出异常或收集失败信息
            }
        }
        return allMilvusData;
    }


    /**
     * Upsert 数据到 Milvus
     * @param milvusDataTuples
     * @param client
     */
    private void upsertDataToMilvus(List<MilvusDataTuple> milvusDataTuples, MilvusServiceClient client) {
        if (milvusDataTuples.isEmpty()) {
            log.info("持久化服务 (Milvus Upsert): 没有数据需要 Upsert 到 Milvus。");
            return;
        }
        log.info("持久化服务 (Milvus Upsert): 成功生成 {} 条数据，准备 Upsert 到 Milvus。", milvusDataTuples.size());

        List<Long> mysqlIds = milvusDataTuples.stream().map(MilvusDataTuple::mysqlId).collect(Collectors.toList());
        List<List<Float>> vectors = milvusDataTuples.stream().map(MilvusDataTuple::vector).collect(Collectors.toList());
        List<String> textContents = milvusDataTuples.stream().map(MilvusDataTuple::textContent).collect(Collectors.toList());

        List<UpsertParam.Field> fieldsData = new ArrayList<>();
        fieldsData.add(new UpsertParam.Field(MILVUS_PK_FIELD_NAME, mysqlIds));
        fieldsData.add(new UpsertParam.Field(VECTOR_FIELD_NAME, vectors));
        fieldsData.add(new UpsertParam.Field(TEXT_CONTENT_FIELD_NAME, textContents));

        UpsertParam upsertParam = UpsertParam.newBuilder()
                .withCollectionName(collectionName)
                .withFields(fieldsData)
                .build();

        try {
            R<MutationResult> upsertResponse = client.upsert(upsertParam);
            if (upsertResponse.getStatus() == R.Status.Success.getCode()) {
                MutationResult mutationResult = upsertResponse.getData();
                log.info("持久化服务 (Milvus Upsert): 成功将 {} 条数据 Upsert 到 Milvus。Upserted PKs (部分示例): {}, Success Count: {}, Upsert Count: {}",
                        mysqlIds.size(),
                        mutationResult.getIDs().getIntId().getDataList().stream().limit(Math.min(5, mysqlIds.size())).collect(Collectors.toList()),
                        mutationResult.getSuccIndexCount(),
                        mutationResult.getUpsertCnt());
            } else {
                log.error("持久化服务 (Milvus Upsert): Upsert 到 Milvus 失败: Status={}, Message={}",
                        upsertResponse.getStatus(), upsertResponse.getMessage());
                if (upsertResponse.getException() != null) {
                    log.error("持久化服务 (Milvus Upsert): Milvus Upsert 异常详情:", upsertResponse.getException());
                }
            }
        } catch (Exception e) {
            log.error("持久化服务 (Milvus Upsert): 执行 Milvus Upsert 操作时出错: {}", e.getMessage(), e);
        }
    }


    /**
     * 校验知识条目
     * @param entry 知识条目
     * @param index 条目索引，如果为 -1，则表示单个条目，否则为批量请求中的索引
     */
    private void validateKnowledgeEntry(KnowledgeEntry entry, int index) {
        String context = (index == -1) ? "单个条目" : "批量请求中第 " + (index + 1) + " 个条目";
        if (entry.getMatchKey() == null || entry.getMatchKey().isBlank()) {
            throw new IllegalArgumentException(context + "必须包含有效的 MatchKey。");
        }
        if (entry.getContent() == null || entry.getContent().isBlank()) {
            throw new IllegalArgumentException(context + " (MatchKey: " + entry.getMatchKey() + ") 内容不能为空。");
        }
    }

    private void validateKnowledgeEntries(List<KnowledgeEntry> entries) {
        for (int i = 0; i < entries.size(); i++) {
            validateKnowledgeEntry(entries.get(i), i);
        }
    }

    private record MilvusDataTuple(Long mysqlId, List<Float> vector, String textContent) {}
}